# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: xguo

import collections
import logging
import json

from absl import app, flags

from coin.experimental.xguo.fruit.base.event_clock import (
    event_clock_scheduler,
    event_clock,
)
from coin.experimental.xguo.fruit.base.strategy_sim_base import (
    BaseSimStrategy,
    BaseStrategyConfig,
)

from coin.experimental.xguo.fruit.util.util import FillStats
from coin.strategy.mm.simple_sim.executors import PassiveSimExecutor
from coin.experimental.xguo.fruit.util.fee import get_fee_rate
from coin.strategy.mm.subscription import SubscriptionRequest
from coin.experimental.xguo.fruit.smm.smm_pricer3 import SmmPricer, SmmPricerConfig
from coin.base.book.util import (
    convert_book_to_dict,
)

FLAGS = flags.FLAGS


class SmmSimStrategy(BaseSimStrategy):
  def __init__(self, config):
    super().__init__(config)
    self.last_fire_time = 0
    self.pnl_report = []
    self._pricers = collections.defaultdict(dict)
    self._fill_stats = collections.defaultdict(lambda: None)
    sub_req = self.strategy_config.strategy_core_config['trade_subreq']
    self._trade_mea = SubscriptionRequest.from_str(sub_req)
    event_clock_scheduler.add_periodic_task(1800, self.log_sim_time)
    event_clock_scheduler.add_periodic_task(300, self.pnl_snapshot)
    event_clock_scheduler.add_periodic_task(3600, self.save_sim_result)
    self._renew_tick_done = False

    for product in self.product_list:
      self.prepare_fill_stat(product)
      self.prepare_executor(product)
    self.prepare_pricers()

  def prepare_fill_stat(self, product):
    if product.exchange != self._trade_mea.exchange:
      self._logger.info('Ignore fill_stat for product %s', product)
      return

    fee_bp = get_fee_rate(product)[0]
    self._logger.info("product %s: fee_bp %s", product, fee_bp)
    self._fill_stats[product] = FillStats(product, fee_bp=fee_bp)
    event_clock_scheduler.add_periodic_task(
        period=30,
        job_func=self._fill_stats[product].print,
        msg='Sim PNL from fill',
    )

  def prepare_pricers(self):
    for product in self.product_list:
      base = product.base.currency
      if self._trade_mea.exchange == product.exchange:
        self._pricers[base]['trade_symbol'] = product
      else:
        self._pricers[base]['ref_symbol'] = product

    pricer_configs = self.strategy_config.strategy_core_config['pricer']
    for base, pricer_info in self._pricers.items():
      trade_symbol = pricer_info['trade_symbol']
      ref_symbol = pricer_info.get('ref_symbol')  # Can be None
      executor = self.executors.get(trade_symbol)
      exe_config = self.strategy_config.executor_config[trade_symbol]['proto_config']
      pricer_config = pricer_configs[trade_symbol.base.currency]
      config = SmmPricerConfig()
      config.trade_symbol = trade_symbol
      config.ref_symbol = ref_symbol
      config.get_pos = lambda: executor.pos.pos
      config.max_pos = exe_config.max_pos
      config.min_pos = exe_config.min_pos
      config.price_range = pricer_config['price_range']

      self._pricers[base]['pricer'] = SmmPricer(config)
      self._logger.info('Add smm pricer for %s', base)

  def prepare_executor(self, product):
    if product.exchange != self._trade_mea.exchange:
      self._logger.info('Ignore executor for product %s', product)
      return

    maker_fee_bp, taker_fee_bp = get_fee_rate(product)
    self._logger.info("product %s fee_bp %s, %s", product, maker_fee_bp, taker_fee_bp)
    config = self.strategy_config.executor_config[product]['proto_config']
    min_pos = config.min_pos
    max_pos = config.max_pos
    lot_size = config.lot_size
    min_fire_period = config.order_update_period_sec * 1e9
    executor = PassiveSimExecutor(
        lot_size=lot_size,
        min_pos=min_pos,
        max_pos=max_pos,
        maker_fee=maker_fee_bp * 1e-4,
        taker_fee=taker_fee_bp * 1e-4,
        execution_delay=100,
        post_only=False,
        min_fire_period=min_fire_period,
        do_agg=True,
    )

    # noinspection PyUnusedLocal
    def sim_fill_callback(price, qty, side, maker, ts):
      self._fill_stats[product].update_sim(price, qty, side)
      fill_dict = {
          'product': str(product),
          'type': 'fill',
          'fill_price': price,
          'fill_qty': qty,
          'side': side,
      }
      self._writer.write(fill_dict, ts)

    executor.fill_callback = sim_fill_callback
    self.executors[product] = executor

  def get_pricer(self, product):
    return self._pricers[product.base.currency]['pricer']

  def on_book(self, product, book):
    timestamp = book.timestamp
    pricer = self.get_pricer(product)
    pricer.update(product, book)
    self._writer.write(convert_book_to_dict(book, product=product), timestamp)

    if product.exchange == self._trade_mea.exchange:
      self._fill_stats[product].update_book(book)
      self.executors[product].update_book(book)

    if not pricer.is_ready():
      return

    if self._renew_tick_done is False:
      for pricer_info in self._pricers.values():
        pricer_info['pricer'].update_tick()
      self._renew_tick_done = True

    if product == pricer.ref_symbol:
      executor = self.executors.get(pricer.trade_symbol)
      assert executor is not None
      orders = pricer.gen_order()
      if orders is None:
        self._logger.info('book is not valid!')
        return
      agg_buy_order = orders.get('agg_buy')
      agg_sell_order = orders.get('agg_sell')
      mm_order = orders.get('mm')
      orders['type'] = 'order'
      self._writer.write(orders, timestamp)

      if agg_buy_order:
        executor.agg_buy(agg_buy_order['price'], timestamp)

      if agg_sell_order:
        executor.agg_sell(agg_sell_order['price'], timestamp)

      if mm_order:
        executor.post(
            timestamp=timestamp,
            buy_price=mm_order['pass_buy_price'],
            sell_price=mm_order['pass_sell_price'])

      self._logger.debug('Submit order: %s', event_clock.now())

  def on_trade(self, product, trade_event):
    executor = self.executors.get(product)
    if executor:
      executor.update_trade(trade_event)

  def pnl_snapshot(self):
    timestamp = event_clock.timestamp
    result = {
        'timestamp': timestamp,
        'human_readable_timestamp': str(event_clock.now()),
        'data': {},
    }
    for product, executor in self.executors.items():
      result['data'][product.symbol] = executor.get_summary()
    print(json.dumps(result, indent=2))
    self.pnl_report.append(result)

  def save_sim_result(self):
    filename = self.strategy_config.sim_config.get('result_file')
    if not filename:
      filename = str(event_clock.timestamp) + '.json'
    with open(filename, 'w') as out_file:
      json.dump(self.pnl_report, out_file)


def main(_):
  config = BaseStrategyConfig.from_file(FLAGS.config)
  strategy = SmmSimStrategy(config)
  strategy.run()
  strategy.save_sim_result()


if __name__ == '__main__':
  flags.DEFINE_string('config', None, 'Strategy config file.')

  logging.basicConfig(level='INFO', format='%(levelname)s %(asctime)s %(name)s] %(message)s')
  app.run(main)
